from contextvars import Token

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request

from src.contexts import clean_user_context, set_user_context
from src.security import verify_token


class UserAuthMiddleware(BaseHTTPMiddleware):
    """获取token并解析后存储到用户上下文中"""

    async def dispatch(self, request: Request, call_next):
        auth_token = request.headers.get("Authorization")
        context_token: Token | None = None
        try:
            if auth_token:
                member_id = verify_token(auth_token.replace("Bearer ", ""))
                context_token = await set_user_context(member_id)
            return await call_next(request)
        finally:
            if context_token is not None:
                clean_user_context(context_token)
